import tensorflow as tf
import os
from PIL import Image
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split

os.environ["CUDA_VISIBLE_DEVICES"] = '-1'


def get_df(DataPath):
    dirs = []
    images = []
    masks = []
    for dirname, _, filenames in os.walk(DataPath):
        for filename in filenames:
            if 'mask' in filename:
                dirs.append(dirname.replace(DataPath, ''))
                masks.append(filename)
                images.append(filename.replace('_mask', ''))
    df = pd.DataFrame({'directory': dirs, 'masks': masks, 'images': images})
    return df


# The following functions can be used to convert a value to a type compatible
# with tf.Example.

def _bytes_feature(value):
    """Returns a bytes_list from a string / byte."""
    if isinstance(value, type(tf.constant(0))):
        value = value.numpy()  # BytesList won't unpack a string from an EagerTensor.
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))


def _float_feature(value):
    """Returns a float_list from a float / double."""
    return tf.train.Feature(float_list=tf.train.FloatList(value=[value]))


def _int64_feature(value):
    """Returns an int64_list from a bool / enum / int / uint."""
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))


def create_example(image_string, mask):
    feature = {
        'mask': _bytes_feature(mask),
        'image': _bytes_feature(image_string),
    }

    return tf.train.Example(features=tf.train.Features(feature=feature))


def generate(imagelist, masks, record_file):
    """使用PIllow 压缩，转numpy的bytes"""

    writer = tf.io.TFRecordWriter(record_file)
    num = 0
    for img, mask in zip(imagelist, masks):
        try:
            image_string = Image.open(img)
            image_string = image_string.convert('RGB')
            image_string = np.asarray(image_string, np.uint8)
            image_string = image_string.tobytes()

            """masks"""
            mask = Image.open(mask).convert('L')
            mask = np.asarray(mask, np.uint8)
            mask = mask.tobytes()
            tf_example = create_example(image_string, mask)
            writer.write(tf_example.SerializeToString())
            num += 1
            if num % 100 == 0:
                print("has create {} example".format(num))
        except Exception as e:
            pass


def read_tf_record(serialized_example):
    image_feature_description = {
        'mask': tf.io.FixedLenFeature([], tf.string),
        'image': tf.io.FixedLenFeature([], tf.string),
    }

    feature = tf.io.parse_example(serialized_example,
                                  features=image_feature_description)
    image = tf.io.decode_raw(feature['image'], tf.uint8)
    image = tf.reshape(image, [256, 256, 3])
    # image = tf.reverse(image, axis=[-1])  # RGB->BGR

    mask = tf.io.decode_raw(feature['mask'], tf.uint8)
    mask = tf.reshape(mask, [256, 256, 1])

    return image, mask


if __name__ == '__main__':
    DataPath = 'E:\\2022Project\\data\\archive\\kaggle_3m\\'
    img_df = get_df(DataPath)
    print(img_df.head())
    shape = img_df.shape[0]

    imagels = []
    maskls = []
    for idx in range(shape):
        imagePath = os.path.join(DataPath, img_df['directory'].iloc[idx], img_df['images'].iloc[idx])
        maskPath = os.path.join(DataPath, img_df['directory'].iloc[idx], img_df['masks'].iloc[idx])
        imagels.append(imagePath)
        maskls.append(maskPath)

    X_train, X_test, y_train, y_test = train_test_split(imagels, maskls, test_size=0.2, random_state=42)

    generate(X_train, y_train, 'train.record')

    generate(X_test, y_test, 'test.record')
